package com.example.graduatedesigned.arima;

import org.apache.commons.math3.stat.regression.OLSMultipleLinearRegression;
import java.text.DecimalFormat;

public class ARIMA {

    public static double[] forecast(double[] data, int p, int d, int q, int forecastSteps) {
        // 对数据进行差分
        double[] diffData = difference(data, d);

        // 训练ARIMA模型
        double[] arCoefficients = trainARModel(diffData, p);
        double[] maCoefficients = trainMAModel(diffData, q);

        // 预测差分后的数据
        double[] forecastedDiff = forecastARIMA(diffData, arCoefficients, maCoefficients, forecastSteps);

        // 反差分得到最终预测结果
        return inverseDifference(forecastedDiff, data, d);
    }

    private static double[] difference(double[] data, int d) {
        double[] result = data.clone();
        for (int i = 0; i < d; i++) {
            double[] temp = new double[result.length - 1];
            for (int j = 0; j < temp.length; j++) {
                temp[j] = result[j + 1] - result[j];
            }
            result = temp;
        }
        return result;
    }

    private static double[] trainARModel(double[] data, int p) {
        int n = data.length;
        if (n <= p) {
            return new double[0];
        }

        OLSMultipleLinearRegression regression = new OLSMultipleLinearRegression();
        double[][] x = new double[n - p][p];
        double[] y = new double[n - p];

        for (int i = 0; i < n - p; i++) {
            for (int j = 0; j < p; j++) {
                x[i][j] = data[i + p - j - 1];
            }
            y[i] = data[i + p];
        }

        regression.newSampleData(y, x);
        return regression.estimateRegressionParameters();
    }

    private static double[] trainMAModel(double[] data, int q) {
        // MA模型训练逻辑，简化实现
        return new double[q];
    }

    private static double[] forecastARIMA(double[] data, double[] arCoefficients, double[] maCoefficients, int steps) {
        int p = arCoefficients.length - 1; // 忽略截距项
        int q = maCoefficients.length;
        int n = data.length;

        double[] forecast = new double[steps];
        double[] combined = new double[n + steps];

        System.arraycopy(data, 0, combined, 0, n);

        for (int i = 0; i < steps; i++) {
            double value = 0;

            // 添加截距项
            if (arCoefficients.length > 0) {
                value += arCoefficients[0];
            }

            // AR部分
            for (int j = 1; j <= p; j++) {
                if (i + n - j >= 0) {
                    value += arCoefficients[j] * combined[i + n - j];
                }
            }

            // MA部分 (简化实现)
            for (int j = 0; j < q; j++) {
                if (i - j - 1 >= 0) {
                    value += maCoefficients[j] * forecast[i - j - 1];
                }
            }

            forecast[i] = value;
            combined[n + i] = value;
        }

        return forecast;
    }

    private static double[] inverseDifference(double[] diffData, double[] originalData, int d) {
        double[] result = new double[diffData.length + d];

        // 复制原始数据的最后d个值作为基础
        System.arraycopy(originalData, originalData.length - d, result, 0, d);
        DecimalFormat df = new DecimalFormat("#.0");

        // 累加差分结果
        for (int i = 0; i < diffData.length; i++) {
            double temp = result[i + d - 1] + diffData[i];
            result[i + d] = Double.parseDouble(df.format(temp));
        }

        return result;
    }
}